import torch
import torch.nn as nn
import torch.optim as optim

# 定义自注意力模型
class SelfAttention(nn.Module):
    def __init__(self, embed_size, heads):
        super(SelfAttention, self).__init__()
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads

        assert self.head_dim * heads == embed_size, "Embed size needs to be divisible by heads"

        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

    def forward(self, values, mask):
        N = values.size(0)
        value_len, key_len, query_len = values.size(1), self.embed_size, self.embed_size

        # Split the embedding into self.heads different pieces
        values = values.reshape(N, value_len, self.heads, self.head_dim)
        keys = self.keys(values).reshape(N, key_len, self.heads, self.head_dim)
        queries = self.queries(values).reshape(N, query_len, self.heads, self.head_dim)

        energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])  # [N, heads, query_len, key_len]

        if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))

        attention = torch.softmax(energy / (self.embed_size ** (1 / 2)), dim=3)
        out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape(N, query_len, self.heads * self.head_dim)
        out = self.fc_out(out)
        return out

# 定义模型参数
d_model = 3  # 维度
heads = 1  # 头数

# 实例化模型
model = SelfAttention(d_model, heads)

# 模拟文本数据和掩码
text = "我是谁"  # 这里需要将文本转换为词嵌入
# 假设我们已经有了词嵌入，这里用随机数模拟
values = torch.randn(1, len(text), d_model)  # [batch_size, seq_len, d_model]

# 假设没有掩码
mask = None

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练一轮
optimizer.zero_grad()
output = model(values, mask)  # 前向传播
loss = criterion(output, torch.tensor([0]))  # 假设标签是0
loss.backward()  # 反向传播
optimizer.step()  # 参数更新

# 打印注意力分数矩阵
attention_scores = model.queries.weight  # 这里只是示例，实际应从模型的输出中获取
print(attention_scores)

if __name__=='__main__':
    print("end")